package cn.net.dpi.shp.util;

import cn.hutool.core.io.FileUtil;
import org.geotools.data.DataStore;
import org.geotools.data.DataStoreFinder;
import org.geotools.data.postgis.PostgisNGDataStoreFactory;
import org.geotools.data.simple.SimpleFeatureSource;
import org.geotools.feature.FeatureCollection;
import org.geotools.feature.FeatureIterator;
import org.geotools.feature.simple.SimpleFeatureTypeBuilder;
import org.jetbrains.annotations.NotNull;
import org.locationtech.jts.geom.MultiPoint;
import org.locationtech.jts.geom.Point;
import org.opengis.feature.simple.SimpleFeature;
import org.opengis.feature.simple.SimpleFeatureType;
import org.opengis.feature.type.AttributeDescriptor;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.IOException;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.PreparedStatement;
import java.sql.Types;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;

public class ShpToPg {

    private static final Logger logger = LoggerFactory.getLogger(ShpToPg.class);

    private static final String PG_HOST = "192.168.200.54";
    private static final String PG_PORT = "5432";
    private static final String PG_DATABASE = "test";
    private static final String PG_USER = "postgres";
    private static final String PG_PASSWORD = "pstgrs@2025";
    private static final String PG_SCHEMA = "public";

    public static void main(String[] args) {
        Arrays.stream(FileUtil.ls("E:\\工作\\项目\\预研\\1d_gps"))
                .filter(f -> f.getName().endsWith(".shp"))
                .forEach(f -> {
                    try {
                        Result result = createTable(f);
                        writeTable(result);
                    } catch (Exception e) {
                        logger.error("file 写入pg失败. {}", e.getMessage(), e);
                    }
                });
    }

    /**
     * 写入数据
     *
     * @param result
     * @throws Exception
     */
    private static void writeTable(Result result) throws Exception {
        FeatureCollection<SimpleFeatureType, SimpleFeature> features = result.featureSource.getFeatures();
        try (FeatureIterator<SimpleFeature> iterator = features.features();
             Connection conn = DriverManager.getConnection(
                     "jdbc:postgresql://" + PG_HOST + ":" + PG_PORT + "/" + PG_DATABASE,
                     PG_USER, PG_PASSWORD)) {
            String insertSQL = buildInsertSQL(result.newSchema);
            PreparedStatement pstmt = conn.prepareStatement(insertSQL);
            int batchSize = 500;
            int count = 0;
            while (iterator.hasNext()) {
                SimpleFeature feature = iterator.next();
                setParameters(pstmt, feature);
                pstmt.addBatch();
                count++;
                if (count % batchSize == 0) {
                    pstmt.executeBatch();
                    logger.info("已提交 {} 条记录", count);
                }
            }
            // 提交剩余未满批次的记录
            if (count % batchSize != 0) {
                pstmt.executeBatch();
                logger.info("最后提交 {} 条记录", (count % batchSize));
            }
            logger.info("数据导入完成");
        }
    }

    /**
     * 建表
     *
     * @param shpFile
     * @return
     * @throws IOException
     */
    @NotNull
    private static Result createTable(File shpFile) throws IOException {
        Map<String, Object> shpParams = new HashMap<>();
        shpParams.put("url", shpFile.toURI().toURL());
        DataStore shpStore = DataStoreFinder.getDataStore(shpParams);

        // 2. 连接PostgreSQL数据库
        Map<String, Object> pgParams = new HashMap<>();
        pgParams.put(PostgisNGDataStoreFactory.DBTYPE.key, "postgis");
        pgParams.put(PostgisNGDataStoreFactory.HOST.key, PG_HOST);
        pgParams.put(PostgisNGDataStoreFactory.PORT.key, PG_PORT);
        pgParams.put(PostgisNGDataStoreFactory.DATABASE.key, PG_DATABASE);
        pgParams.put(PostgisNGDataStoreFactory.SCHEMA.key, PG_SCHEMA);
        pgParams.put(PostgisNGDataStoreFactory.USER.key, PG_USER);
        pgParams.put(PostgisNGDataStoreFactory.PASSWD.key, PG_PASSWORD);
        DataStore pgStore = DataStoreFinder.getDataStore(pgParams);

        // 3. 转换数据
        String typeName = shpStore.getTypeNames()[0];
        SimpleFeatureSource featureSource = shpStore.getFeatureSource(typeName);

        // 4. 创建小写字段名的新Schema
        SimpleFeatureType schema = featureSource.getSchema();
        SimpleFeatureTypeBuilder builder = new SimpleFeatureTypeBuilder();
        String tableName = typeName.toLowerCase();
        builder.setName(tableName);
        for (AttributeDescriptor attr : schema.getAttributeDescriptors()) {
            builder.add(attr.getLocalName().toLowerCase(), attr.getType().getBinding());
        }
        SimpleFeatureType newSchema = builder.buildFeatureType();
        // 5. 创建小写表名的新表（自动创建表）
        try {
            pgStore.createSchema(newSchema);
        } catch (Exception e) {
            String message = e.getCause().getMessage();
            if (message != null && message.contains("already exists")) {
                logger.info("table {} already exists", tableName);
            } else {
                throw e;
            }
        }
        pgStore.getFeatureWriterAppend(tableName, null).close();
        // 6. 关闭连接
        shpStore.dispose();
        pgStore.dispose();
        return new Result(featureSource, newSchema);
    }

    private static String buildInsertSQL(SimpleFeatureType schema) {
        StringBuilder sql = new StringBuilder("INSERT INTO " + schema.getTypeName().toLowerCase() + " (");
        StringBuilder values = new StringBuilder(" VALUES (");

        for (int i = 0; i < schema.getAttributeCount(); i++) {
            if (i > 0) {
                sql.append(", ");
                values.append(", ");
            }
            sql.append(schema.getDescriptor(i).getLocalName());
            values.append("?");
        }

        sql.append(")").append(values).append(")");
        return sql.toString();
    }

    private static void setParameters(PreparedStatement pstmt, SimpleFeature feature) throws Exception {
        for (int i = 0; i < feature.getAttributeCount(); i++) {
            Object o = feature.getAttribute(i);
            if (o instanceof Point) {
                pstmt.setObject(i + 1, o, Types.OTHER);
            } else if (o instanceof MultiPoint) {
                MultiPoint theGeom = (MultiPoint) o;
                Point geometry = (Point) theGeom.getGeometryN(0);
                pstmt.setObject(i + 1, geometry, Types.OTHER);
            } else {
                pstmt.setObject(i + 1, o);
            }

        }
    }

    private static class Result {
        public final SimpleFeatureSource featureSource;
        public final SimpleFeatureType newSchema;

        public Result(SimpleFeatureSource featureSource, SimpleFeatureType newSchema) {
            this.featureSource = featureSource;
            this.newSchema = newSchema;
        }
    }
}
